from jmxquery import JMXConnection, JMXQuery

# 1. 要开启JMX 2. 要能识别到Java

kafka_jmx = "172.16.100.74:9999"
MyJmxConnection = JMXConnection(f"service:jmx:rmi:///jndi/rmi://{kafka_jmx}/jmxrmi")


def print_metric(_type_str: str, _attribute_value: str, _attribute_list: list):
    if _type_str == "Memory":
        bean_name_str = "java.lang:type={0}".format(_type_str)
        for _name in _attribute_list:
            jmx_query = [JMXQuery(mBeanName=bean_name_str, attribute=_name,
                                  metric_name="{type}_{attribute}_{attributeKey}")]
            _metrics = MyJmxConnection.query(jmx_query)
            for _metric in _metrics:
                print(f"{_metric.to_query_string()} ({_metric.value_type}) = {_metric.value}")
    elif _type_str == "MemoryPool":
        for _name in _attribute_list:
            bean_name_str = f"java.lang:type={_type_str},name={_name}"
            jmx_query = [JMXQuery(mBeanName=bean_name_str, attribute=_attribute_value,
                                  metric_name="{type}_{attribute}_{attributeKey}")]
            _metrics = MyJmxConnection.query(jmx_query)
            for _metric in _metrics:
                print(f"{_metric.to_query_string()} ({_metric.value_type}) = {_metric.value}")
    elif _type_str == "KafkaController":
        for _name in _attribute_list:
            bean_name_str = f"kafka.controller:type={_type_str},name={_name}"
            jmx_query = [JMXQuery(mBeanName=bean_name_str, attribute=_attribute_value,
                                  metric_name="{type}_{attribute}_{attributeKey}")]
            _metrics = MyJmxConnection.query(jmx_query)
            for _metric in _metrics:
                print(f"{_metric.to_query_string()} ({_metric.value_type}) = {_metric.value}")
    elif _type_str == "SocketServer":
        for _name in _attribute_list:
            bean_name_str = f"kafka.network:type={_type_str},name={_name}"
            jmx_query = [JMXQuery(mBeanName=bean_name_str, attribute=_attribute_value,
                                  metric_name="{type}_{attribute}_{attributeKey}")]
            _metrics = MyJmxConnection.query(jmx_query)
            for _metric in _metrics:
                print(f"{_metric.to_query_string()} ({_metric.value_type}) = {_metric.value}")
    else:
        print("传入的类型不对")


# 服务端指标
# 集群活跃controller个数：ActiveControllerCount
controller_list = ["ControllerState", "ActiveControllerCount"]
print_metric("KafkaController", "Value", controller_list)
# broker网络空闲率：NetworkProcessorAvgIdlePercent
network_list = ["NetworkProcessorAvgIdlePercent"]
print_metric("SocketServer", "Value", network_list)
# memory
memory_list = ["HeapMemoryUsage", "NonHeapMemoryUsage"]
print_metric("Memory", "None", memory_list)
# memoryPool
memoryPool_list = ["Code Cache", "Metaspace", "Compressed Class Space", "Par Eden Space", "Par Survivor Space",
                   "PS Eden Space", "PS Old Gen", "PS Perm Gen", "PS Survivor Space", "CMS Old Gen", "CMS Perm Gen",
                   "Perm Gen"]
print_metric("MemoryPool", "Usage", memoryPool_list)

## GarbageCollector 其中时间单位是s
type_str = "GarbageCollector"
for name in ["Copy", "MarkSweepCompact", "PS Scavenge", "ConcurrentMarkSweep", "ParNew", "PS MarkSweep"]:
    beanName_str = "java.lang:type={0},name={1}".format(type_str, name)
    for atrributeName in ["CollectionTime", "CollectionCount"]:
        jmxQuery = [
            JMXQuery(mBeanName=beanName_str,
                     attribute=atrributeName,
                     metric_name="Type:{type}_{attribute}")]
        metrics = MyJmxConnection.query(jmxQuery)
        for metric in metrics:
            print(f"{metric.to_query_string()} ({metric.value_type}) = {metric.value}")

## Threading
jmxQuery = [
    JMXQuery(mBeanName="java.lang:type=Threading",
             attribute="ThreadCount",
             metric_name="Type:{type}_{attribute}")]
metrics = MyJmxConnection.query(jmxQuery)
for metric in metrics:
    print(f"{metric.to_query_string()} ({metric.value_type}) = {metric.value}")
